#include "Grammar.h"
#include <iostream>

void Grammar::addProduction(const Symbol &left, const Production &right)
{
    productions_[left].push_back(right);
}

void Grammar::setTerminals(const std::unordered_set<Symbol> &terminals)
{
    terminals_ = terminals;
}

bool Grammar::isTerminal(const Symbol &symbol) const
{
    // 约定：不在产生式左侧的符号且在终结符集合中的是终结符
    if (terminals_.count(symbol))
        return true;
    if (productions_.count(symbol) == 0)
        return true; // 没有产生式定义的视为终结符
    return false;
}

void Grammar::computeFirst()
{
    // 对每个符号调用递归计算FIRST
    for (const auto &p : productions_)
    {
        const Symbol &nonterminal = p.first;
        std::unordered_set<Symbol> visited;
        computeFirstForSymbol(nonterminal, visited);
    }
}

// 递归计算某符号FIRST集，visited防止递归死循环
bool Grammar::computeFirstForSymbol(const Symbol &symbol, std::unordered_set<Symbol> &visited)
{
    // 如果已经计算过了，跳过
    if (first_.count(symbol) && !first_[symbol].empty())
    {
        return first_[symbol].count("epsilon") > 0;
    }

    // 避免循环调用
    if (visited.count(symbol))
    {
        return false;
    }
    visited.insert(symbol);

    // 终结符的FIRST集是自己
    if (isTerminal(symbol))
    {
        first_[symbol].insert(symbol);
        return false;
    }

    // 非终结符，遍历产生式
    bool canProduceEpsilon = false;
    for (const auto &prod : productions_[symbol])
    {
        bool allCanEpsilon = true;

        for (const auto &sym : prod)
        {
            // 递归计算子符号FIRST集
            bool subCanEpsilon = computeFirstForSymbol(sym, visited);

            // 把子符号的FIRST（除epsilon）加入当前FIRST
            for (const auto &s : first_[sym])
            {
                if (s != "epsilon")
                    first_[symbol].insert(s);
            }

            if (!subCanEpsilon)
            {
                allCanEpsilon = false;
                break;
            }
        }

        if (allCanEpsilon)
        {
            first_[symbol].insert("epsilon");
            canProduceEpsilon = true;
        }
    }

    return canProduceEpsilon;
}

const std::unordered_set<Grammar::Symbol> &Grammar::getFirst(const Symbol &symbol) const
{
    static std::unordered_set<Symbol> emptySet;
    auto it = first_.find(symbol);
    if (it != first_.end())
        return it->second;
    return emptySet;
}

void Grammar::printFirst() const
{
    std::cout << "FIRST集合:\n";
    for (const auto &p : first_)
    {
        std::cout << "FIRST(" << p.first << ") = { ";
        bool first = true;
        for (const auto &sym : p.second)
        {
            if (!first)
                std::cout << ", ";
            std::cout << sym;
            first = false;
        }
        std::cout << " }\n";
    }
}

void Grammar::computeFollow(const Symbol &startSymbol)
{
    // 初始化：将 $ 放入起始符号的 FOLLOW 集
    follow_[startSymbol].insert("$");

    bool changed;
    do
    {
        changed = false;
        for (const auto &prodPair : productions_)
        {
            const Symbol &A = prodPair.first;
            for (const auto &prod : prodPair.second)
            {
                for (size_t i = 0; i < prod.size(); ++i)
                {
                    const Symbol &B = prod[i];
                    if (isTerminal(B))
                        continue; // 终结符没有 FOLLOW 集

                    // B 后面的符号
                    bool allCanEpsilon = true;
                    for (size_t j = i + 1; j < prod.size(); ++j)
                    {
                        const Symbol &beta = prod[j];
                        if (isTerminal(beta))
                        {
                            // 终结符直接加入 FOLLOW(B)
                            if (follow_[B].insert(beta).second)
                            {
                                changed = true;
                            }
                            allCanEpsilon = false;
                            break;
                        }

                        // 把 FIRST(beta)-epsilon 加入 FOLLOW(B)
                        for (const auto &f : first_[beta])
                        {
                            if (f != "epsilon")
                            {
                                if (follow_[B].insert(f).second)
                                {
                                    changed = true;
                                }
                            }
                        }
                        if (first_[beta].count("epsilon") == 0)
                        {
                            allCanEpsilon = false;
                            break;
                        }
                    }

                    // 如果 B 后面全能 epsilon，或者 B 已在最后，把 FOLLOW(A) 加入 FOLLOW(B)
                    if (allCanEpsilon || i + 1 == prod.size())
                    {
                        for (const auto &f : follow_[A])
                        {
                            if (follow_[B].insert(f).second)
                            {
                                changed = true;
                            }
                        }
                    }
                }
            }
        }
    } while (changed);
}

void Grammar::printFollow() const
{
    std::cout << "FOLLOW集合:\n";
    for (const auto &p : follow_)
    {
        std::cout << "FOLLOW(" << p.first << ") = { ";
        bool first = true;
        for (const auto &sym : p.second)
        {
            if (!first)
                std::cout << ", ";
            std::cout << sym;
            first = false;
        }
        std::cout << " }\n";
    }
}

const std::unordered_set<Grammar::Symbol> &Grammar::getFollow(const Symbol &symbol) const
{
    static std::unordered_set<Symbol> emptySet;
    auto it = follow_.find(symbol);
    if (it != follow_.end())
        return it->second;
    return emptySet;
}

const std::unordered_map<Grammar::Symbol, std::vector<Grammar::Production>> &Grammar::getProductions() const
{
    return productions_;
}
